import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split


x, y = load_digits(return_X_y=True)
x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=1000, shuffle=False)

bytes = 0

with open("datasets/digits.bin", "wb") as f:
    bytes += len(x_train.tobytes())
    bytes += len(np.eye(10)[y_train].tobytes())

    print(len(x_train.tobytes()))
    print(len(np.eye(10)[y_train].tobytes()))

    bytes += len(x_test.tobytes())
    bytes += len(np.eye(10)[y_test].tobytes())

    print(len(x_test.tobytes()))
    print(len(np.eye(10)[y_test].tobytes()))

    f.write(x_train.tobytes())
    f.write(x_test.tobytes())
    f.write(np.eye(10)[y_train].tobytes())
    f.write(np.eye(10)[y_test].tobytes())

print(f"bytes {bytes}")
